import requests
import os

api_key = os.getenv("BAILIAN_QWEN_API_KEY")

class InterLMM:
    def __init__(self):
        self.base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
        self.headers = {
            "Content-Type": "application/json",
            f"Authorization": f"Bearer {api_key}"
        }
        self.params = {
            "temperature": 0.1,
            "top_p": 0.001,
            "model": "deepseek-r1-distill-llama-70b",
            "chat_template_kwargs": {"enable_thinking": False},
            "messages": [{"role": "user", "content": "你是谁？"}],
            "stream": False
        }
        # self.base_url = "http://127.0.0.1:8999/v1/chat/completions"
        # self.params = {
        #     "temperature": 0.1,
        #     "top_p": 0.001,
        #     "model": "Qwen",
        #     "chat_template_kwargs": {"enable_thinking": False},
        #     "messages": [{"role": "user", "content": "你是谁？"}],
        #     "stream": False
        # }

    def ask_single(self, question):
        self.params["messages"][0]["content"] = question
        response = self.send_request()
        if response:
            return response['choices'][0]['message']['content']
        else:
            return None

    def send_request(self):
        try:
            # response = requests.post(self.base_url, json=self.params)
            response = requests.post(self.base_url, json=self.params, headers=self.headers)
            response.raise_for_status()  # 如果状态码不是200，会抛出HTTPError异常
            return response.json()
        except requests.RequestException as e:
            print(f"请求失败: {e}")
            return None


# 使用示例
if __name__ == "__main__":
    interlmm = InterLMM()
    question = """
从机关、单位名称中提取下列信息：
1）去掉地域名称的机关、单位名称。
   例如：克州农村合作经济发展指导中心， 去掉地域后名称为 农村合作经济发展指导中心
2）机构性质：党委、政府和两院、事业单位、人大、政协、团体、企业
（人员检察院、人民法院、纪委监察机关算入政府和两院）

- 输入信息：
  1）机构名称：博尔塔拉蒙古自治州行政服务中心
  2）机构所在地域：乌鲁木齐市水磨沟区

- 输出格式：json，不输出json以外的文字、信息；json字段包括 org_name（机构名称）、nature（机构性质）。示例如下

{
 "org_name":"机构名称",
 "nature":"机构性质"
}
"""
    response = interlmm.ask_single(question)
    print("响应：", response)
